D:\a\csshw\csshw\src\lib.rs
Line | Count | Source |
1 | | //! Cluster SSH tool for Windows inspired by csshX |
2 | | |
3 | | #![deny(clippy::implicit_return)] |
4 | | #![allow(clippy::needless_return, clippy::doc_overindented_list_items)] |
5 | | #![warn(missing_docs)] |
6 | | #![doc(html_no_source)] |
7 | | |
8 | | use std::fs::{create_dir, File}; |
9 | | use std::mem; |
10 | | |
11 | | use log::warn; |
12 | | use registry::{value, Data, Hive, Security}; |
13 | | use simplelog::{format_description, ConfigBuilder, LevelFilter, WriteLogger}; |
14 | | use windows::core::PWSTR; |
15 | | use windows::Win32::Foundation::HWND; |
16 | | use windows::Win32::System::Threading::{PROCESS_INFORMATION, STARTUPINFOW}; |
17 | | |
18 | | #[cfg(test)] |
19 | | use mockall::automock; |
20 | | |
21 | | pub mod cli; |
22 | | pub mod client; |
23 | | pub mod daemon; |
24 | | pub mod serde; |
25 | | pub mod utils; |
26 | | |
27 | | use utils::windows::WindowsApi; |
28 | | |
29 | | /// CLSID identifying `conhost.exe` in the registry. |
30 | | /// |
31 | | /// As used in Windows Terminal: |
32 | | /// <https://github.com/microsoft/terminal/blob/v1.22.3232.0/src/propslib/DelegationConfig.hpp#L105> |
33 | | const CLSID_CONHOST: &str = "{B23D10C0-E52E-411E-9D5B-C09FDF709C7D}"; |
34 | | /// CLSID identifying the default configuration in the registry. |
35 | | /// |
36 | | /// The default configuration is "let windows choose". |
37 | | /// Also defined in Windows Terminal: |
38 | | /// <https://github.com/microsoft/terminal/blob/v1.22.3232.0/src/propslib/DelegationConfig.hpp#L104> |
39 | | const CLSID_DEFAULT: &str = "{00000000-0000-0000-0000-000000000000}"; |
40 | | /// Registry path where `DelegationConsole` and `DelegationTerminal` registry keys are stored. |
41 | | /// |
42 | | /// These registry keys store the configuration value for the default terminal application. |
43 | | const DEFAULT_TERMINAL_APP_REGISTRY_PATH: &str = r"Console\%%Startup"; |
44 | | /// `DelegationConsole` registry key. |
45 | | /// |
46 | | /// As used in Windows Terminal: |
47 | | /// <https://github.com/microsoft/terminal/blob/v1.22.3232.0/src/propslib/DelegationConfig.cpp#L29> |
48 | | const DELEGATION_CONSOLE: &str = "DelegationConsole"; |
49 | | /// `DelegationTerminal` registry key. |
50 | | /// |
51 | | /// As used in Windows Terminal: |
52 | | /// <https://github.com/microsoft/terminal/blob/v1.22.3232.0/src/propslib/DelegationConfig.cpp#L30> |
53 | | const DELEGATION_TERMINAL: &str = "DelegationTerminal"; |
54 | | |
55 | | /// Trait for registry operations to enable mocking in tests |
56 | | #[cfg_attr(test, automock)] |
57 | | pub trait Registry { |
58 | | /// Get a string value from the registry |
59 | | fn get_registry_string_value(&self, path: &str, name: &str) -> Option<String>; |
60 | | /// Set a string value in the registry |
61 | | fn set_registry_string_value(&self, path: &str, name: &str, value: &str) -> bool; |
62 | | } |
63 | | |
64 | | /// Default implementation of Registry trait that performs actual Windows registry API calls |
65 | | pub struct DefaultRegistry; |
66 | | |
67 | | impl Registry for DefaultRegistry { |
68 | 0 | fn get_registry_string_value(&self, path: &str, name: &str) -> Option<String> { |
69 | 0 | let key = Hive::CurrentUser |
70 | 0 | .open(path, Security::Read | Security::Write) |
71 | 0 | .ok()?; |
72 | 0 | match key.value(name) { |
73 | 0 | Ok(Data::String(value)) => return Some(value.to_string_lossy()), |
74 | 0 | Ok(_) => panic!("Expected string data for {name} registry value"), |
75 | 0 | Err(value::Error::NotFound(_, _)) => return Some(CLSID_DEFAULT.to_owned()), |
76 | 0 | Err(err) => { |
77 | 0 | warn!("Failed to read {} value from registry: {}", name, err); |
78 | 0 | return None; |
79 | | } |
80 | | } |
81 | 0 | } |
82 | | |
83 | 0 | fn set_registry_string_value(&self, path: &str, name: &str, value: &str) -> bool { |
84 | 0 | if let Ok(key) = Hive::CurrentUser.open(path, Security::Read | Security::Write) { |
85 | 0 | match key.set_value::<String>( |
86 | 0 | name.to_owned(), |
87 | 0 | &Data::String(value.to_owned().try_into().unwrap()), |
88 | 0 | ) { |
89 | 0 | Ok(()) => return true, |
90 | | Err(_) => { |
91 | 0 | warn!("Failed to set registry value {} to {}", name, value); |
92 | 0 | return false; |
93 | | } |
94 | | } |
95 | | } else { |
96 | 0 | return false; |
97 | | } |
98 | 0 | } |
99 | | } |
100 | | |
101 | | /// Return the Window Handle [HWND] for the foreground window associated with the given `process_id`. |
102 | | /// |
103 | | /// If multiple foreground windows are associated with the given `process_id` it is undefined which [HWND] gets returned. |
104 | | /// |
105 | | /// # Arguments |
106 | | /// |
107 | | /// * `windows_api` - Windows API operations implementation |
108 | | /// * `process_id` - ID of the process for which to retrieve the window handle. |
109 | | /// |
110 | | /// # Returns |
111 | | /// |
112 | | /// The Window Handle [HWND] for the window associated with the given `process_id`. |
113 | 0 | pub fn get_console_window_handle<W: WindowsApi>(windows_api: &W, process_id: u32) -> HWND { |
114 | 0 | return windows_api.get_window_handle_for_process(process_id); |
115 | 0 | } |
116 | | |
117 | | /// Create process with command line using the provided API (testable version) |
118 | | /// |
119 | | /// # Arguments |
120 | | /// |
121 | | /// * `api` - Windows API operations implementation |
122 | | /// * `application` - Application name including file extension |
123 | | /// * `command_line` - UTF-16 encoded command line |
124 | | /// |
125 | | /// # Returns |
126 | | /// |
127 | | /// [PROCESS_INFORMATION] of the spawned process or None if failed |
128 | 3 | pub fn create_process<W: WindowsApi>( |
129 | 3 | api: &W, |
130 | 3 | application: &str, |
131 | 3 | command_line: &[u16], |
132 | 3 | ) -> Option<PROCESS_INFORMATION> { |
133 | 3 | let mut startupinfo = STARTUPINFOW { |
134 | 3 | cb: mem::size_of::<STARTUPINFOW>() as u32, |
135 | 3 | ..Default::default() |
136 | 3 | }; |
137 | 3 | let mut process_information = PROCESS_INFORMATION::default(); |
138 | 3 | let mut cmd_line = command_line.to_vec(); |
139 | 3 | let command_line_ptr = PWSTR(cmd_line.as_mut_ptr()); |
140 | | |
141 | 3 | match api.create_process_raw( |
142 | 3 | application, |
143 | 3 | command_line_ptr, |
144 | 3 | &mut startupinfo, |
145 | 3 | &mut process_information, |
146 | 3 | ) { |
147 | 2 | Ok(()) => return Some(process_information), |
148 | 1 | Err(_) => return None, |
149 | | } |
150 | 3 | } |
151 | | |
152 | | /// Trait for file system operations to enable mocking in tests |
153 | | #[cfg_attr(test, automock)] |
154 | | pub trait FileSystem { |
155 | | /// Create a directory |
156 | | fn create_directory(&self, path: &str) -> bool; |
157 | | /// Create a log file |
158 | | fn create_log_file(&self, filename: &str) -> bool; |
159 | | } |
160 | | |
161 | | /// Default implementation of FileSystem trait that performs actual file system operations |
162 | | pub struct ProductionFileSystem; |
163 | | |
164 | | impl FileSystem for ProductionFileSystem { |
165 | 0 | fn create_directory(&self, path: &str) -> bool { |
166 | 0 | return create_dir(path).is_ok() || std::path::Path::new(path).exists(); |
167 | 0 | } |
168 | | |
169 | 0 | fn create_log_file(&self, filename: &str) -> bool { |
170 | 0 | return File::create(filename).is_ok(); |
171 | 0 | } |
172 | | } |
173 | | |
174 | | /// Guard storing previous/old `DelegationConsole` and `DelegationTerminal` registry values. |
175 | | /// |
176 | | /// Configures `conhost.exe` as the default terminal application |
177 | | /// and reverts to the original configuration when being dropped. |
178 | | pub struct WindowsSettingsDefaultTerminalApplicationGuard<R: Registry> { |
179 | | /// Old `DelegationConsole` registry value |
180 | | old_windows_terminal_console: Option<String>, |
181 | | /// Old `DelegationTerminal` registry value |
182 | | old_windows_terminal_terminal: Option<String>, |
183 | | /// Registry operations trait |
184 | | registry: R, |
185 | | } |
186 | | |
187 | | impl<R: Registry> WindowsSettingsDefaultTerminalApplicationGuard<R> { |
188 | | /// Create a new guard with the given registry operations |
189 | | /// |
190 | | /// # Arguments |
191 | | /// |
192 | | /// * `registry` - Registry operations implementation |
193 | | /// |
194 | | /// # Returns |
195 | | /// |
196 | | /// A new guard that will restore registry values on drop |
197 | 5 | pub fn new_with_registry(registry: R) -> Self { |
198 | 5 | let mut guard = WindowsSettingsDefaultTerminalApplicationGuard { |
199 | 5 | old_windows_terminal_console: None, |
200 | 5 | old_windows_terminal_terminal: None, |
201 | 5 | registry, |
202 | 5 | }; |
203 | | |
204 | 3 | if let (Some(console_val), Some(terminal_val)) = ( |
205 | 5 | guard |
206 | 5 | .registry |
207 | 5 | .get_registry_string_value(DEFAULT_TERMINAL_APP_REGISTRY_PATH, DELEGATION_CONSOLE), |
208 | 5 | guard |
209 | 5 | .registry |
210 | 5 | .get_registry_string_value(DEFAULT_TERMINAL_APP_REGISTRY_PATH, DELEGATION_TERMINAL), |
211 | | ) { |
212 | | // No need to change if already set to conhost |
213 | 3 | if console_val == CLSID_CONHOST && terminal_val == CLSID_CONHOST1 { |
214 | 1 | return guard; |
215 | 2 | } |
216 | | |
217 | | // Store old values and set new ones |
218 | 2 | guard.old_windows_terminal_console = Some(console_val); |
219 | 2 | guard.old_windows_terminal_terminal = Some(terminal_val); |
220 | | |
221 | 2 | guard.registry.set_registry_string_value( |
222 | 2 | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
223 | 2 | DELEGATION_CONSOLE, |
224 | 2 | CLSID_CONHOST, |
225 | | ); |
226 | 2 | guard.registry.set_registry_string_value( |
227 | 2 | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
228 | 2 | DELEGATION_TERMINAL, |
229 | 2 | CLSID_CONHOST, |
230 | | ); |
231 | | } else { |
232 | 2 | warn!( |
233 | 0 | "Failed to read registry key {}, \ |
234 | 0 | cannot make sure conhost.exe is the configured default terminal application", |
235 | | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
236 | | ); |
237 | | } |
238 | | |
239 | 4 | return guard; |
240 | 5 | } |
241 | | } |
242 | | |
243 | | impl WindowsSettingsDefaultTerminalApplicationGuard<DefaultRegistry> { |
244 | | /// Create a new guard with production registry operations |
245 | 0 | pub fn new() -> Self { |
246 | 0 | return Self::new_with_registry(DefaultRegistry); |
247 | 0 | } |
248 | | } |
249 | | |
250 | | impl<R: Registry> Default for WindowsSettingsDefaultTerminalApplicationGuard<R> |
251 | | where |
252 | | R: Default, |
253 | | { |
254 | 0 | fn default() -> Self { |
255 | 0 | return Self::new_with_registry(R::default()); |
256 | 0 | } |
257 | | } |
258 | | |
259 | | impl Default for DefaultRegistry { |
260 | 0 | fn default() -> Self { |
261 | 0 | return DefaultRegistry; |
262 | 0 | } |
263 | | } |
264 | | |
265 | | impl<R: Registry> Drop for WindowsSettingsDefaultTerminalApplicationGuard<R> { |
266 | | /// Restore the original default terminal application setting to the registry. |
267 | | /// |
268 | | /// If old values weren't stored, nothing is done. |
269 | 5 | fn drop(&mut self) { |
270 | 2 | if let (Some(old_console), Some(old_terminal)) = ( |
271 | 5 | &self.old_windows_terminal_console, |
272 | 5 | &self.old_windows_terminal_terminal, |
273 | 2 | ) { |
274 | 2 | self.registry.set_registry_string_value( |
275 | 2 | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
276 | 2 | DELEGATION_CONSOLE, |
277 | 2 | old_console, |
278 | 2 | ); |
279 | 2 | self.registry.set_registry_string_value( |
280 | 2 | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
281 | 2 | DELEGATION_TERMINAL, |
282 | 2 | old_terminal, |
283 | 2 | ); |
284 | 3 | } |
285 | 5 | } |
286 | | } |
287 | | |
288 | | /// Launch the given console application with the given arguments as a new detached process with its own console window. |
289 | | /// |
290 | | /// Input/Output handles are not being inherited. |
291 | | /// Whichever default terminal application is configured in the windows system settings will be used |
292 | | /// to host the application (i.e. create the window). |
293 | | /// |
294 | | /// # Arguments |
295 | | /// |
296 | | /// * `api` - Windows API implementation |
297 | | /// * `application` - Application name including file extension (`.exe`). |
298 | | /// If the application is not in the `PATH` environment variable, the full path |
299 | | /// must be specified. |
300 | | /// * `args` - List of arguments to the application. |
301 | | /// |
302 | | /// # Returns |
303 | | /// |
304 | | /// [PROCESS_INFORMATION] of the spawned process. |
305 | 4 | pub fn spawn_console_process<W: WindowsApi>( |
306 | 4 | api: &W, |
307 | 4 | application: &str, |
308 | 4 | args: Vec<String>, |
309 | 4 | ) -> Option<PROCESS_INFORMATION> { |
310 | 4 | return api.create_process_with_args(application, args); |
311 | 4 | } |
312 | | |
313 | | /// Initialize the logger. |
314 | | /// |
315 | | /// Makes sure a `logs` directory exists in the current working directory. |
316 | | /// Log filename format: `<utc-time-of-executable-start>_<name>.log`. |
317 | | /// Configures [log_panics]. |
318 | | /// |
319 | | /// # Arguments |
320 | | /// |
321 | | /// * `name` - Will be part of the log filename. |
322 | 0 | pub fn init_logger(name: &str) { |
323 | 0 | init_logger_with_fs(&ProductionFileSystem, name); |
324 | 0 | } |
325 | | |
326 | | /// Initialize the logger with the provided file system operations. |
327 | | /// |
328 | | /// # Arguments |
329 | | /// |
330 | | /// * `fs` - File system operations implementation |
331 | | /// * `name` - Will be part of the log filename |
332 | 9 | pub fn init_logger_with_fs<F: FileSystem>(fs: &F, name: &str) { |
333 | 9 | let utc_now = chrono::offset::Utc::now() |
334 | 9 | .format("%Y-%m-%d_%H-%M-%S.%f") |
335 | 9 | .to_string(); |
336 | | |
337 | 9 | fs.create_directory("logs"); |
338 | | |
339 | 9 | let filename = format!("logs/{utc_now}_{name}.log"); |
340 | 9 | if fs.create_log_file(&filename) { |
341 | 7 | if let Ok(file0 ) = File::create(&filename) { |
342 | 0 | let _ = WriteLogger::init( |
343 | 0 | LevelFilter::Debug, |
344 | 0 | ConfigBuilder::new() |
345 | 0 | .set_time_format_custom(format_description!( |
346 | 0 | "[hour]:[minute]:[second].[subsecond]" |
347 | 0 | )) |
348 | 0 | .build(), |
349 | 0 | file, |
350 | 0 | ); |
351 | 0 | log_panics::init(); |
352 | 7 | } |
353 | 2 | } |
354 | 9 | } |
355 | | |
356 | | /// Detect if application was launched from Windows Explorer (GUI) vs command line using the provided console API. |
357 | | /// |
358 | | /// Returns true if launched from GUI (separate console), false if from existing console. |
359 | | /// Based on: <https://stackoverflow.com/a/513574> |
360 | | /// |
361 | | /// # Arguments |
362 | | /// |
363 | | /// * `windows_api` - Windows API operations implementation |
364 | | /// |
365 | | /// # Returns |
366 | | /// |
367 | | /// * `true` - Application was launched from GUI (Explorer, double-click, etc.) |
368 | | /// * `false` - Application was launched from existing console (command line) |
369 | 7 | pub fn is_launched_from_gui<W: WindowsApi>(windows_api: &W) -> bool { |
370 | 7 | match windows_api.get_stdout_handle() { |
371 | 6 | Ok(handle) => { |
372 | 6 | match windows_api.get_console_screen_buffer_info_with_handle(handle) { |
373 | 5 | Ok(csbi) => { |
374 | | // The cursor has not moved from the initial 0,0 position -> launched in separate console |
375 | 5 | return csbi.dwCursorPosition.X == 0 && csbi.dwCursorPosition.Y == 01 ; |
376 | | } |
377 | 1 | Err(err) => { |
378 | 1 | warn!("GetConsoleScreenBufferInfo failed: {:?}"0 , err); |
379 | 1 | return false; |
380 | | } |
381 | | } |
382 | | } |
383 | 1 | Err(err) => { |
384 | 1 | warn!("Failed to get stdout handle: {:?}"0 , err); |
385 | 1 | return false; |
386 | | } |
387 | | } |
388 | 7 | } |
389 | | |
390 | | #[cfg(test)] |
391 | | #[path = "./tests/test_lib.rs"] |
392 | | mod test_lib; |